
from gensim.models import Word2Vec  # 词向量
from random import choice
from os.path import exists
from warnings import filterwarnings
filterwarnings('ignore')  # 不打印警告


class CONFIG:
    path = 'model/古诗词.txt'
    window = 16  # 滑窗大小
    min_count = 60  # 过滤低频字
    size = 125  # 词向量维度
    topn = 14  # 生成诗词的开放度
    model_path = 'model/word2vec'


class Model:
    def __init__(self, window, topn, model):
        self.window = window
        self.topn = topn
        self.model = model  # 词向量模型
        self.chr_dict = model.wv.index_to_key  # 字典

    @classmethod
    def initialize(cls, config):
        """模型初始化"""
        if exists(config.model_path):
            # 模型读取
            model = Word2Vec.load(config.model_path)
        else:
            # 语料读取
            with open(config.path, encoding='utf-8') as f:
                ls_of_ls_of_c = [list(line.strip()) for line in f]
            # 模型训练和保存
            model = Word2Vec(ls_of_ls_of_c, window=config.window, min_count=config.min_count)
            model.save(config.model_path)
        return cls(config.window, config.topn, model)

    def poem_generator(self, keyword, form):
        """古诗词生成"""
        def clean(list):
            return [t[0] for t in list if t[0] not in ['，', '。']]
        # # 标题补全
        # if len(title) < 4:
        #     if not title:
        #         title += choice(self.chr_dict)
        #     for _ in range(4 - len(title)):
        #         similar_chr = self.model.wv.similar_by_key(title[-1], self.topn // 2)
        #         similar_chr = clean(similar_chr)
        #         nextChar = choice([ for  in similar_chr if  not in title])
        #         title += nextChar
        # 文本生成
        poem = list(keyword)
        for i in range(form[0]):
            for _ in range(form[1]):
                predict_char = self.model.predict_output_word(
                    poem[-self.window:], max(self.topn, len(poem) + 1))
                predict_char = clean(predict_char)
                nextChar = choice([char for char in predict_char if char not in poem[len(keyword):]])
                poem.append(nextChar)
            poem.append('，' if i % 2 == 0 else '。\r')
        length = form[0] * (form[1] + 1)
        return ''.join(poem[-length:])


def main(keyword,config=CONFIG):
    """主函数"""
    form = {'五言绝句': (4, 5), '七言绝句': (4, 7),'五言律诗':(8, 5),'七言律诗':(8, 7)}
    model = Model.initialize(config)
    while True:
        random = choice([1, 2, 3, 4])
        # keyword = input('输入关键词：').strip()
        if random == 1:
            poem1 = model.poem_generator(keyword, form['五言绝句'])
            print(poem1)
            return poem1
        if random == 2:
            poem2 = model.poem_generator(keyword, form['七言绝句'])
            print(poem2)
            return poem2
        if random == 3:
            poem3 = model.poem_generator(keyword, form['五言律诗'])
            print(poem3)
            return poem3
        if random == 4:
            poem4 = model.poem_generator(keyword, form['七言律诗'])
            print(poem4)
            return poem4



if __name__ == '__main__':
    main()